 #!/usr/bin/env python
#Filename: atlasgen.py

import math
import cPickle
import itertools as it
from mpl_toolkits.mplot3d import Axes3D
import matplotlib.pyplot as pp
from numpy import linalg as LA
import imageTools.ITKUtils.io as io
import numpy as np
import random as rand

def generate_atlas(graphlist):

	#calculate maxrank of all graphs analyzed
	maxrank=0
	maxsize = (0,512,512)
	for i in range(len(graphlist)):

		#unpack
		gfname = graphlist[i]
		gfo = file(gfname, 'rb')
		graph = cPickle.load(gfo)
		gfo.close()

		og = graph['orderedGraphs'][(1,'mp_graphs')]
		nodes = og.nodes(data=True)
		edges = og.edges(data=True)

		"""
		#for each graph examine the image size, keep largest dataframe
		if graph['imgShape'][0]>maxsize[0]:	maxsize = graph['imgShape']	
		"""

		#find maxrank of nodes
		#print nodes
		for n in nodes:
			print n
			if n[1]['rank']>maxrank: maxrank = n[1]['rank']
		print "done with nodes...maxrank=",maxrank
		#check for rank
		#print "graph nodes before dumping...",graph['orderedGraphs'][(1,'mp_graphs')].nodes(True)
		gfo = file(gfname,'wb')
		cPickle.dump(graph,gfo)
		gfo.close()
	
	#build ranking-tally dictionary:
	#each rank contains a dictionary of the mean, variance and coordinate list
	ranking = {}
	for r in range(maxrank+1):
		ranking['{0}'.format(r)] = {'mean': [], 'var': [], 'coord': []}

	#load up ranking dictionary with node coordinates
	for i in range(len(graphlist)):
		gfname = graphlist[i]
		gfo = file(gfname, 'rb')
		data = cPickle.load(gfo)
		og = data['orderedGraphs'][(1,'mp_graphs')]
		nodes = og.nodes(data=True)
	
		#for a given graph, search all nodes one rank at a time
		for rank in range(maxrank+1):
			for n in nodes:
				#print "search graph for rank:",graphlist[i]
				#print "check node:",n
				try:
					#if node matches current rank, add it to the coordinate list
					if n[1]['rank']==rank:  ranking['{0}'.format(rank)]['coord'].append(n[1]['trwcrd'])
				except: pass
		#finished with file
		gfo.close()

	#calculate mean and variance coordinates
	meancrd = lambda x: map(np.mean,x)		#returns [meanx's,meany's,meanz's]
	varcrd = lambda x: map(np.var,x)		#returns [varx's,vary's,varz's]
	for rank in range(maxrank+1):
		x = np.array(ranking['{0}'.format(rank)]['coord'])	#convert coordinates into nparray
		#print "points", x
		x = x.transpose()					#[xyz1,xyz2...]--->[x's,y's,z's]
		#print "coordinates", x
		ranking['{0}'.format(rank)]['mean'] = meancrd(x)
		ranking['{0}'.format(rank)]['var'] = varcrd(x)

	atlas = {'ranking':ranking, 'graphlist': graphlist}

	return atlas


def autorotate(graph,atlas,iterate=None):
	"""rotates graph in atlas scalarfield until membership is maximized"""
	#parse graph data
	og = graph['orderedGraphs'][(1,'mp_graphs')]
	root = graph['roots'][(1,'mp_graphs')]
	nodes = og.nodes(data=True)
	edges = og.edges(data=True)
	#rootwcrd = nodes[og.nodes().index(root)][1]['trwcrd']

	#initialize
	bestang = [0,0,0]
	bestscore = graph_score(graph,atlas,closest_node)
 	#print "start:", bestscore
	if not iterate == None:
		iterations = iterate
	else:
		iterations=500
	print "Optimizing rotations..."
	#MAXIMIZE MEMBERSHIP SCORE: find bestang and bestscore
	for N in range(iterations):

		#generate angles
		angles = random_angles(bestang[0],bestang[1],bestang[2])
		
		#rotate
		edge = []
		edge = rotate_edges(np.radians(angles),edges,'d0t')
		#print "rotated output:",edge
		
		#save in graph
		i=0			
		for e in edge:			
			og.edges(data=True)[i][2]['d0t'] = e
			#print og.edges(data=True)[i][2].keys()
			i = i+1

		#calculate membership score
		score = graph_score(graph,atlas,closest_node)
		#print "iteration:",N," score:",score

		if score>bestscore:
			bestscore = score
			bestang = angles
			#print "iteration:", N, "bestscore: {0}  bestang: {1}".format(bestscore,bestang)
			
		else:
			continue

	print "Done optimizing! :DDD"	

	#Apply best rotation after iterations
	print "bestscore: {0}  bestang: {1}".format(bestscore,bestang)
	bestedges = rotate_edges(np.radians(bestang),edges,'d0t')
	membership = bestscore	
	trwcrdnodes = rotate_nodes(np.radians(bestang), graph)
	#plotgraph(graph)

	#load final transformed edges into net-x edge dict	(overwrite iterations)
	i=0			
	for e in bestedges:			
		og.edges(data=True)[i][2]['d0t'] = e
		i = i+1
	#load final transformed nodes into networkx nodes dictionary	(overwrite)		
	i=0			
	for n in trwcrdnodes:
		og.nodes(data=True)[i][1]['trwcrd'] = n
		i=i+1
	
	return graph

def center_graph(graph,imgname):

	img = io.readImage(imgname,imgMode='sshort')
	scale = img.GetSpacing()

	#parse graph data
	og = graph['orderedGraphs'][(1,'mp_graphs')]
	root = graph['roots'][(1,'mp_graphs')]
	nodes = og.nodes(data=True)
	edges = og.edges(data=True)
	rootwcrd = nodes[og.nodes().index(root)][1]['wcrd']

	#CENTER/SCALE GRAPH

	shift = np.ndarray(shape=(1,3))
	shift = rootwcrd
	
	#-----NODES------
	#convert
	if nodes == og.nodes(data = True):
		for n in range(len(nodes)):		#tuple ---> list, tuple ---> array
			nodes[n] = list(nodes[n])
			nodes[n][0] = np.array(nodes[n][0])
	#rank nodes
	nodes = rank_nodes(graph)	

	#load world coordinates
	wcrdnodes = []
	for n in nodes:
		wcrdnodes.append(n[1]['wcrd'])

	#recenter nodes by rootwcrd
	trwcrdnodes = []
	for n in wcrdnodes:
		trwcrdnodes.append(n-shift)

	#rescale nodes
	for n in trwcrdnodes:
	
		n[0] = n[0]*scale[0]	
		n[1] = n[1]*scale[1]
		n[2] = n[2]*scale[2]
	
	#store transformed nodes in networkx nodes dictionary		
	i=0			
	for n in trwcrdnodes:
		og.nodes(data=True)[i][1]['trwcrd'] = n
		i=i+1

	#-----EDGES-----
	#load edge list as array elements
	edges = []
	for e in og.edges(data=True):
		try:
			edges.append(np.array(e[2]['d0']))	#e[2]['d0'] = array([x's,y's,z's])
		except:
			pass

	#recenter edges around rootwcrd
	for i in range(len(edges)):

		edges[i][0] = edges[i][0] - shift[0]		#recenter x's by shift[x]
		edges[i][1] = edges[i][1] - shift[1]		#y's
		edges[i][2] = edges[i][2] - shift[2]		#z's
	
	#rescale graph
	for i in range(len(edges)):

		#rescale edge components to mm-scale
		edges[i][0] = edges[i][0]*scale[0]		#x's
		edges[i][1] = edges[i][1]*scale[1]		#y's
		edges[i][2] = edges[i][2]*scale[2]		#z's
	
	#store transformed edges in net-x edges dict
	i=0			
	for e in edges:			
		og.edges(data=True)[i][2]['d0t'] = e
		i = i+1

	return graph


def rank_nodes(graph):
    """accepts networkx dictionary, returns same modified dictionary
    use adjacency list to trace nodes through different levels of rank"""
    print "ranking nodes..."
    #unpack
    og = graph['orderedGraphs'][(1,'mp_graphs')]
    root = graph['roots'][(1,'mp_graphs')]
    nodes = og.nodes(data=True)
    edges = og.edges(data=True)		
    adjac = og.adjacency_list()
    rank = 0
    
    #assign rank=0 to root node in nodes()
    nodes[og.nodes(data=False).index(root)][1]['rank']= 0
    
    for rank in range(len(adjac)):
        indices =[]
        curns =[]		#current nodes
        #load current nodes of 'rank'==rank, mark their index location in nodes list
        for n in nodes:
            try:
                if n[1]['rank']==rank: 
                    curns.append(n[0])
                    indices.append(nodes.index(n))
            except:
                pass
        #load a complete list of children nodes from adjacency list
        childs=[]
        for i in indices:
            for c in adjac[i]:
                childs.append(c)
    
        #search for each child node in nodes list; update child node 'rank' = rank+1
        for c in childs:
            for n in nodes:
                if n[0]==c: n[1]['rank']=rank+1
    print "...done ranking nodes"
    return nodes


def graph_score(graph,atlas,fxn,order=None):
    """uses an atlas's spatial function to quantify the graph's similarity to the atlas"""
    #unpack
    ranking = atlas['ranking']
    og = graph['orderedGraphs'][(1,'mp_graphs')]
    edges = og.edges(data=True)
    
    total = 0
    length = 0
    
    #trace each edge of graph, calculate it's path integral
    for e in edges:
        
        if order == None:
            #load coordinates
            try:				#try transformed points	
                xs = e[2]['d0t'][0]
                ys = e[2]['d0t'][1]
                zs = e[2]['d0t'][2]
                points = list(it.izip(xs,ys,zs))
                #print "graphscore points:", points
            except:				#else use regular edge data
                xs = e[2]['d0'][0]
                ys = e[2]['d0'][1]
                zs = e[2]['d0'][2]
                points = list(it.izip(xs,ys,zs))
        else:
                xs = e[2][str(order)][0]
                ys = e[2][str(order)][1]
                zs = e[2][str(order)][2]
                points = list(it.izip(xs,ys,zs))
                
        #path integral of edge's centerline
        # should this be changed to cython?
        for i in range(len(points)): total = total + fxn(points[i],atlas)
        length = length + len(points)
    
    #score is cumulative since total and length are cumulative
    score = total/length
    return score

def score_by_rank(graph,atlas):

	og = graph['orderedGraphs'][(1,'mp_graphs')]
	nodes = og.nodes(data=True)

	total=0
	length = len(nodes)

	for n in nodes:
		point = n[1]['trwcrd']
		#print "node:",n,"point:",point
		rank = n[1]['rank']
		#print "point:",point, "rank:",rank
		gausslist = gauss_nodes(point,atlas)
		#print gausslist
		try:
			scalar = gausslist[rank]
			total = total + scalar
			#if not scalar == None: print "chose:",scalar,"total =",total
		except:
			scalar = closest_node(point,atlas)
			total = total + scalar
			#if not scalar == None: #print "chose:",scalar,"total =",total
	
	score = total/length
	return score



def closest_node(point,atlas):
    """returns the mean node to which the point has closest variance
    list gaussianscores of sequential node rank"""
    gausslist = gauss_nodes(point,atlas)
    maxscore = max(gausslist)
    rank = gausslist.index(maxscore)	#since gausslist is organized by ascending rank
    mean = atlas['ranking']['{0}'.format(rank)]	#retrieve mean coord associated with highest score
    
    return maxscore


def gauss_layering(point,atlas):
    """returns scalar value associated with a point's cumulative
    gaussian relation to a spatial node ranking"""
    gaussset = gauss_nodes(point, atlas)
    
    #normalize by number of ranks
    scalar = sum(gaussset)/len(gaussset)
    return scalar


def gauss_nodes(point,atlas):
    """returns a 3D gaussian value for each rank of node in the atlas
    NOTE: can be used to determine which rank of node a point is closest to"""
    ranking = atlas['ranking']
    weight = 1
    p = {'x':point[0],'y':point[1],'z':point[2]}	
    stats = {}
    
    #1D gaussian
    gauss = lambda (x,mx,vx): math.exp(-(x-mx)**2/(2*vx))
    gaussset= []
    
    #unpack ranking
    for rank in range(len(ranking)):
    
        current = ranking['{0}'.format(rank)]
    
        #stats dict
        
        #divide by zero condition
        for i in range(3):
            if current['var'][i] == 0:
                current['var'][i] = 1
    
        stats['x'] = [current['mean'][0],current['var'][0]]
        stats['y'] = [current['mean'][1],current['var'][1]]
        stats['z'] = [current['mean'][2],current['var'][2]]
        
        #axis gaussians
        gx = gauss((p['x'],stats['x'][0],stats['x'][1]))
        gy = gauss((p['y'],stats['y'][0],stats['y'][1]))
        gz = gauss((p['z'],stats['z'][0],stats['z'][1]))
    
        gauss3d = weight*gx*gy*gz
        gaussset.append(gauss3d)
    
    return gaussset


def map_layer((mini,maxi), atlas, color, window, scale):
    """generates a set of coordinates of points which meet certain minmax conditions
    returns a dictionary containing coordinates, boundary conditions and color-mapping"""
    if maxi == None: maxi=1
    if mini == None: mini=0
    
    layer = { 'color': color, 'bounds': (mini,maxi)}
    
    #fxn = atlas['fxn']
    fxn = closest_node
    
    grid = gen3Dgrid(window,scale)
    layer['coords'] = np.array(filter(lambda p: mini<fxn(p,atlas)<maxi, grid))
    
    return layer

def gen3Dgrid(window, scale):
    """accepts window and scale parameters, returns grid array"""
    
    #generate set of coordinates (R3)
    xs = np.arange(window[0][0],window[0][1],scale)
    ys = np.arange(window[1][0],window[1][1],scale)
    zs = np.arange(window[2][0],window[2][1],scale)
    
    #generate cartesian product of grid points
    total = np.array(list(it.product(xs,ys,zs)))
    
    return total

def plotgraph(graph,fignum=None):
    if not fignum == None:
        fig = pp.figure(fignum)
        ax = Axes3D(fig)
    else:
        fig = pp.figure(1)
        ax = Axes3D(fig)
    
    og = graph['orderedGraphs'][(1,'mp_graphs')]
    edges = og.edges(data=True)
    for e in edges:
        try:			#plot transformed graph if possible
            sp = e[2]['d0t']
            ax.plot(sp[0],sp[1],sp[2])
            pp.show()
        except:			#else try plotting non-transformed
            try:
                sp = e[2]['d0']
                ax.plot(sp[0],sp[1],sp[2])
                pp.show()
            except:	pass		#else graph data is corrupted
    return

def plotatlas(atlas,window,scale, numlayers=None):
	
	red = map_layer((0.7,0.9),atlas,'r', window, scale )	
	#print "red",len(red['coords'].transpose())
	blue = map_layer((0.44,0.65),atlas,'b',window, scale )
	#print "blue",len(blue['coords'].transpose())
	grn = map_layer((0.3,0.4),atlas,'g', window, scale )
	#print "grn",len(grn['coords'].transpose())
	
	#plot grid
	fig1 = pp.figure(1)
	ax = Axes3D(fig1)
	
	try:	
		pr = red['coords'].transpose()
		ax.scatter(pr[2],pr[1],pr[0],c=red['color'],marker = 'o',alpha=0.8)
		pp.show()
	except:	pass
	try:	
		pb = blue['coords'].transpose()
		ax.scatter(pb[2],pb[1],pb[0],c=blue['color'],marker='d',alpha=0.7)
		pp.show()
	except:	pass	
	try:
		pg = grn['coords'].transpose()
		ax.scatter(pg[2],pg[1],pg[0],c=grn['color'],marker='o',alpha=0.5)
		pp.show()
	except:	pass	

	return

def plotatlasandgraph(atlas,graph,window,scale):

	fig = pp.figure(1)
	ax = Axes3D(fig)

	red = map_layer((0.6,0.7),atlas,'r', window, scale )	
	#print "red",len(red['coords'].transpose())
	blue = map_layer((0.44,0.55),atlas,'b',window, scale )
	#print "blue",len(blue['coords'].transpose())
	grn = map_layer((0.3,0.4),atlas,'g', window, scale )
	#print "grn",len(grn['coords'].transpose())
	
	#plot grid
	fig1 = pp.figure(1)
	ax = Axes3D(fig1)
	
	try:	
		pr = red['coords'].transpose()
		ax.scatter(pr[2],pr[1],pr[0],c=red['color'],marker = 'o',alpha=0.8)
		pp.show()
	except:	pass
	try:	
		pb = blue['coords'].transpose()
		ax.scatter(pb[2],pb[1],pb[0],c=blue['color'],marker='d',alpha=0.7)
		pp.show()
	except:	pass	
	try:
		pg = grn['coords'].transpose()
		ax.scatter(pg[2],pg[1],pg[0],c=grn['color'],marker='o',alpha=0.5)
		pp.show()
	except:	pass

	og = graph['orderedGraphs'][(1,'mp_graphs')]
	edges = og.edges(data=True)
	for e in edges:
	    try:			#plot transformed graph if possible
		sp = e[2]['d0t']
		ax.plot(sp[0],sp[1],sp[2])
		pp.show()
	    except:			#else try plotting non-transformed
		try:
			sp = e[2]['d0']
			ax.plot(sp[0],sp[1],sp[2])
			pp.show()
		except:	pass		#else graph data is corrupted
	
	return

#accepts degrees
#distributions are triangular centered around best angles
def random_angles(bestx,besty,bestz):
	angx = rand.triangular(0,360,bestx)
	angy = rand.triangular(0,360,besty)
	angz = rand.triangular(0,360,bestz)
	angles = (angx,angy,angz)
	return angles

#accepts an edge dictionary
#returns rotated set of edges specified by order = 'd0' or 'd1' or 'd2'
def rotate_edges(rads, edges, order=None):
	
	if order == None: order = 'd0'
	
	#zip edges [x's,y's,z's] --> [ [xyz1], [xyz2]... [xyzn] ]
	rows = []
	for e in edges:
		points = np.array(zip(e[2][order][0],e[2][order][1],e[2][order][2]))
		rows.append(points)
		#print rows
	
	#apply rotations
	edge = []
	for r in rows:
		vecs = np.array(map(lambda v: rotation('x',rads[0],v), r))
		vecs = np.array(map(lambda v: rotation('y',rads[1],v), vecs))
		vecs = np.array(map(lambda v: rotation('z',rads[2],v), vecs))
		
		#unzip		
		coords = np.array(zip(*vecs))
		edge.append(coords)

	return edge

def rotate_nodes(rads,graph):
	og = graph['orderedGraphs'][(1,'mp_graphs')]
	chars = ['x','y','z']
	nodes = []
	
	for n in og.nodes(data=True):
		nodes.append(n[1]['trwcrd'])

	for i in range(len(rads)):
		nodes = np.array(map(lambda v: rotation(chars[i],rads[i],v), nodes))
	return nodes

def rotation(axis, rads, vector):

	#rotate around x axis
	Rx = np.matrix([ [1,0,0], [0, np.cos(rads), -np.sin(rads)], [0, np.sin(rads), np.cos(rads)] ])	
	#y axis
	Ry = np.matrix([ [np.cos(rads),0, np.sin(rads)], [0,1,0], [-np.sin(rads),0, np.cos(rads)] ])	
	#z axis
	Rz = np.matrix([ [np.cos(rads), -np.sin(rads),0], [np.sin(rads), np.cos(rads),0], [0,0,1] ])		

	Rot = {'x':Rx, 'y':Ry, 'z':Rz}	#for quick, organized use
	
	try:
		vector = vector.reshape(3,1)	
		vector = Rot[axis]*vector		#type(vector)= np.matrixlib.defmatrix.matrix
		vector = np.array(vector.ravel()[0])[0]	#convert type(vector) == np.ndarray, shape(1,3)

		return vector
	
	except:
		return

#End of atlasgen.py
